#include <cmath>
#include <dirent.h>
#include <string.h>
#include <map>
#include <vector>
#include <chrono>
#include "acllite_dvpp_lite/ImageProc.h"
#include "acllite_om_execute/ModelProc.h"
#include "opencv2/opencv.hpp"

using namespace std;
using namespace acllite;
using namespace cv;

// 定义一个结构体来存储边界框的信息
typedef struct BoundBox {
	float x;
	float y;
	float width;
	float height;
	float score;
	size_t classIndex;
  size_t index;
} BoundBox;

// 定义一个比较函数，用于排序边界框的得分
bool sortScore(BoundBox box1, BoundBox box2)
{
	return box1.score > box2.score;
}

int main(){    
    vector<string> label = {{"hjch"}};
    AclLiteResource aclResource;// 初始化ACL资源 
    bool ret = aclResource.Init();
    CHECK_RET(ret, LOG_PRINT("[ERROR] InitACLResource failed."); return 1);
    
    // 初始化图像处理和模型处理对象
    ImageProc imageProc;
    ModelProc modelProc;
    uint32_t modelWidth = 640;
    uint32_t modelHeight = 640;
    // 加载模型
    ret = modelProc.Load("../model/hjch.om");
    
    CHECK_RET(ret, LOG_PRINT("[ERROR] load model hjch.om failed."); return 1);
    
    // 设置输入图片和输出文件夹的路径
  	string imagesPath = "../data/"; 
    string outputFolderPath = "../output/yolov8/"; 
 
    DIR *dir;// 打开图片文件夹
    struct dirent *ent;
    if ((dir = opendir(imagesPath.c_str())) != NULL) 
    {
        // 遍历文件夹中的所有文件
        while ((ent = readdir(dir)) != NULL) 
        {
            string fileName = ent->d_name;
            
            // 检查文件是否为JPG格式
            if (fileName.find(".jpg") != string::npos) 
            {
                string imagePath = imagesPath + fileName;
    
    ImageData src = imageProc.Read(imagePath);// 读取图片
    CHECK_RET(src.size, LOG_PRINT("[ERROR] ImRead image failed."); return 1);
    printf("%d %d %d\n",src.size,src.width,src.height);
    
    auto start = std::chrono::high_resolution_clock::now();// 记录推理开始时间

    // 创建一个新的图像数据对象用于存储处理后的图像
    ImageData dst;
    ImageSize dsize(640, 640);
    imageProc.Resize(src, dst, dsize,RESIZE_PROPORTIONAL_UPPER_LEFT); // 对图像进行缩放
    printf("%d %d %d\n",dst.size,dst.width,dst.height);

    // 创建模型输入
    ret = modelProc.CreateInput(static_cast<void *>(dst.data.get()), dst.size);
    CHECK_RET(ret, LOG_PRINT("[ERROR] Create model input failed."); return 1);

    // 执行模型推理
    vector<InferenceOutput> inferOutputs;
    ret = modelProc.Execute(inferOutputs);
    CHECK_RET(ret, LOG_PRINT("[ERROR] model execute failed."); return 1);
    
    uint32_t dataSize = inferOutputs[0].size;
    printf("inferOutputs[0].size=%d\n",dataSize);
    // get result from output data set
    // 获取类别概率数据
	uint32_t outputDataBufId = 0;  
    float *classBuff = static_cast<float *>(inferOutputs[outputDataBufId].data.get());   
   
    float confidenceThreshold = 0.25;  // 设置置信度阈值
    size_t classNum = 1;    // 设置类别数量
    size_t offset = 4;   // 设置偏移量
    size_t totalNumber = classNum + offset;  
    size_t modelOutputBoxNum = 8400;    // 设置模型输出的边界框数量
    
    size_t startIndex = 4;   // 设置起始索引
    // 读取原始图片
    cv::Mat srcImage = cv::imread(imagePath);
    int srcWidth = srcImage.cols; 
    int srcHeight = srcImage.rows;

    vector <BoundBox> boxes;  // 初始化边界框列表
    size_t yIndex = 1;  
    size_t widthIndex = 2;  
    size_t heightIndex = 3;  
    size_t classConfidenceIndex = 4; 
    // 计算宽度和高度的缩放比例
    float widthScale = (float)(srcWidth) / modelWidth;
    float heightScale = (float)(srcHeight) / modelHeight;
    float finalScale = (widthScale > heightScale) ? widthScale : heightScale; 
    // 遍历模型输出的边界框
    for (size_t i = 0; i < modelOutputBoxNum; ++i)  
    {
        
        float classScore = classBuff[offset * modelOutputBoxNum + i];// 获取类别得分
        // 如果得分高于置信度阈值，则处理该边界框
        if (classScore > confidenceThreshold)  
        {
            BoundBox box;     
            box.x = classBuff[i] * finalScale; 
            box.y = classBuff[i+modelOutputBoxNum * yIndex] * finalScale;
            box.width = classBuff[i+modelOutputBoxNum * widthIndex] * finalScale;
            box.height = classBuff[i + modelOutputBoxNum * heightIndex] * finalScale;
            box.score = classScore;
            box.classIndex = 0;
            box.index = i;
            boxes.push_back(box);
            
        }
    }
    // 对边界框进行非极大值抑制NMS
    vector <BoundBox> result;  
    result.clear();
    float NMSThreshold = 0.45;   
    int32_t maxLength = modelWidth > modelHeight ? modelWidth : modelHeight;  
    std::sort(boxes.begin(), boxes.end(), sortScore);  // 根据得分对边界框进行排序
    BoundBox boxMax;
    BoundBox boxCompare;    
    while (boxes.size() != 0)  //执行NMS
    {
        size_t index = 1;
        result.push_back(boxes[0]);
        while (boxes.size() > index) {
            boxMax.score = boxes[0].score;
            boxMax.classIndex = boxes[0].classIndex;
            boxMax.index = boxes[0].index;
            boxMax.x = boxes[0].x + maxLength * boxes[0].classIndex;
            boxMax.y = boxes[0].y + maxLength * boxes[0].classIndex;
            boxMax.width = boxes[0].width;
            boxMax.height = boxes[0].height;

            boxCompare.score = boxes[index].score;
            boxCompare.classIndex = boxes[index].classIndex;
            boxCompare.index = boxes[index].index;
            boxCompare.x = boxes[index].x + boxes[index].classIndex * maxLength;
            boxCompare.y = boxes[index].y + boxes[index].classIndex * maxLength;
            boxCompare.width = boxes[index].width;
            boxCompare.height = boxes[index].height;
            float xLeft = max(boxMax.x, boxCompare.x);
            float yTop = max(boxMax.y, boxCompare.y);
            float xRight = min(boxMax.x + boxMax.width, boxCompare.x + boxCompare.width);
            float yBottom = min(boxMax.y + boxMax.height, boxCompare.y + boxCompare.height);
            float width = max(0.0f, xRight - xLeft);
            float hight = max(0.0f, yBottom - yTop);
            float area = width * hight;
            float iou =  area / (boxMax.width * boxMax.height + boxCompare.width * boxCompare.height - area);
            if (iou > NMSThreshold) {
                boxes.erase(boxes.begin() + index);
                continue;
            }
            ++index;
        }
        boxes.erase(boxes.begin());
    }

    const double fountScale = 0.5;
    const uint32_t lineSolid = 2;
    const uint32_t labelOffset = 11;
    const cv::Scalar fountColor(0, 0, 255);
    const vector <cv::Scalar> colors{
        cv::Scalar(237, 149, 100), cv::Scalar(0, 215, 255),
        cv::Scalar(50, 205, 50), cv::Scalar(139, 85, 26)};

    int half = 2;
    string textPrint = "[";
    for (size_t i = 0; i < result.size(); ++i) {
        if (result[i].score < 0.7) {
            continue;
        }
        cv::Point leftUpPoint, rightBottomPoint;
        leftUpPoint.x = result[i].x - result[i].width / half;
        leftUpPoint.y = result[i].y - result[i].height / half;
        rightBottomPoint.x = result[i].x + result[i].width / half;
        rightBottomPoint.y = result[i].y + result[i].height / half;
        cv::rectangle(srcImage, leftUpPoint, rightBottomPoint, colors[i % colors.size()], lineSolid);
        string className = label[result[i].classIndex];
        string markString = to_string(result[i].score) + ":" + className;
        textPrint += markString;
        cv::putText(srcImage, markString, cv::Point(leftUpPoint.x, leftUpPoint.y + labelOffset),
                    cv::FONT_HERSHEY_COMPLEX, fountScale, fountColor);
    }
    textPrint += "]";
                // 记录结束时间
                auto end = std::chrono::high_resolution_clock::now(); 
                std::chrono::duration<float, std::milli> inferenceTime = end - start;
                cout << "Inference time for " << fileName << " is " << inferenceTime.count() << " ms" << endl;
                // 保存处理后的图片
                string savePath = outputFolderPath + fileName;
                cv::imwrite(savePath, srcImage); 
            }
        }
        closedir(dir);// 关闭文件夹
    } else {
        LOG_PRINT("[ERROR] Open directory failed.");
        return 1;
    }

    return 0;
}

